"""
1,检测两个图像的特征点，计算特征描术符(使用SURF)
2,使用FLANN方法进行特征匹配
3,使用得到的匹配点对进行F矩阵计算
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
import ransac
import misc_func as mf

def skew(a):
    """ Skew matrix A such that a x v = Av for any v. """
    return np.array([[0,-a[2],a[1]],[a[2],0,-a[0]],[-a[1],a[0],0]])

"""已测，与opencv中代码结果一致"""
def seven_point_algorithm(X,X_p):
    """此方法参考自opencv源代码中的fundam.cpp中的run7Point函数"""
    A=np.ones([7,9])
    x0=X[:,0]; y0=X[:,1]; x1=X_p[:,0]; y1=X_p[:,1]
    #形成A矩阵
    for i in range(7):
        x0=X[i,0]; y0=X[i,1]; x1=X_p[i,0]; y1=X_p[i,1]
        A[i,0]=x1*x0;        A[i,1]=x1*y0;        A[i,2]=x1
        A[i,3]=y1*x0;        A[i,4]=y1*y0;        A[i,5]=y1
        A[i,6]=x0;           A[i,7]=y0;           A[i,8]=1
    #应用SVD分解，得到f1,f2
    U,S,V=np.linalg.svd(A)  #注意返回的V实际是V^T
    f1=V[7,:]; f2=V[8,:]    #取V的最后两列

    #det(F)=det(aF1-(1-a)F2)=0,得到关于a的三次多项式方程
    c=np.zeros([4],dtype=np.double)     #多项式系数
    f1=f1-f2
    t0 = f2[4]*f2[8] - f2[5]*f2[7]
    t1 = f2[3]*f2[8] - f2[5]*f2[6]
    t2 = f2[3]*f2[7] - f2[4]*f2[6]
    c[3] = f2[0]*t0 - f2[1]*t1 + f2[2]*t2
    c[2] = f1[0]*t0 - f1[1]*t1 + f1[2]*t2 - \
            f1[3]*(f2[1]*f2[8] - f2[2]*f2[7]) + \
            f1[4]*(f2[0]*f2[8] - f2[2]*f2[6]) - \
            f1[5]*(f2[0]*f2[7] - f2[1]*f2[6]) + \
            f1[6]*(f2[1]*f2[5] - f2[2]*f2[4]) - \
            f1[7]*(f2[0]*f2[5] - f2[2]*f2[3]) + \
            f1[8]*(f2[0]*f2[4] - f2[1]*f2[3])
    t0 = f1[4]*f1[8] - f1[5]*f1[7]
    t1 = f1[3]*f1[8] - f1[5]*f1[6]
    t2 = f1[3]*f1[7] - f1[4]*f1[6]
    c[1] = f2[0]*t0 - f2[1]*t1 + f2[2]*t2 - \
            f2[3]*(f1[1]*f1[8] - f1[2]*f1[7]) + \
            f2[4]*(f1[0]*f1[8] - f1[2]*f1[6]) - \
            f2[5]*(f1[0]*f1[7] - f1[1]*f1[6]) + \
            f2[6]*(f1[1]*f1[5] - f1[2]*f1[4]) - \
            f2[7]*(f1[0]*f1[5] - f1[2]*f1[3]) + \
            f2[8]*(f1[0]*f1[4] - f1[1]*f1[3])
    c[0] = f1[0]*t0 - f1[1]*t1 + f1[2]*t2
    n,aa=cv2.solveCubic(c)    #解多项式

    Fs=[]
    for i in range(n):   #xs有三个解，可能是复数（丢弃）
        a=aa[i];    F=np.zeros([9])
        lam =a; mu=1.0
        s=f1[8]*a+f2[8]
        if(np.abs(s) > 2.22E-16):      ##normalize, so that F[3,3]=1
            mu=1/s; lam =lam*mu;    F[8]=1
        else:
            F[8]=0
        for i in range(8):
            F[i]=f1[i]*lam+f2[i]*mu
        F=F.reshape([3,3])
        Fs.append(F)
    return Fs

"""已测，与opencv中代码结果一致"""
def eight_point_algorithm(X,Xp):
    """此方法参考自opencv源代码中的fundam.cpp中的run8Point函数"""
    ##计算中心点以及平均距离
    N=len(X)
    m1c=np.array([0,0])
    m2c=np.array([0,0])
    for i in range(N):
        m1c=m1c+X[i,:]
        m2c=m2c+Xp[i,:]
    m1c=m1c/N;  m2c=m2c/N

    scale1=0; scale2=0
    for i in range(N):
        scale1 +=np.linalg.norm(X[i,:]-m1c)
        scale2 +=np.linalg.norm(Xp[i,:]-m2c)

    scale1=scale1/N; scale2=scale2/N
    scale1 = np.sqrt(2.0)/scale1
    scale2 = np.sqrt(2.0)/scale2

    ##形成一个线性方程Ax=0：每个点对组成一行
    AA=np.zeros([9,9])  #计算A^T*A矩阵，这是一个半正定对称矩阵，所以可用特征值分解
    for i in range(N):
        x0=(X[i,0]-m1c[0])*scale1
        y0=(X[i,1]-m1c[1])*scale1
        x1=(Xp[i,0]-m2c[0])*scale2
        y1=(Xp[i,1]-m2c[1])*scale2
        r=np.array([[x1*x0,x1*y0,x1,y1*x0,y1*y0,y1,x0,y0,1]])
        AA=AA+r.transpose()*r

    #进行对称阵的特征值分解，得到W为特征值list，V为对应特征向量矩阵
    W,V=np.linalg.eigh(AA)

    #取V的第一列为Af的解，因为W是升序排序的
    F0=V[:,0]
    F0=F0.reshape([3,3])

    #对F0进行非满秩化处理
    U,D,V=np.linalg.svd(F0) #D为奇异值list
    D[2]=0
    F=np.dot(U,np.dot(np.diag(D),V))
    #应用反变换，得到真正F
    T1=np.array([scale1,0,-scale1*m1c[0], 0, scale1, -scale1*m1c[1],0,0,1])
    T1=T1.reshape([3,3])
    T2=np.array([scale2,0,-scale2*m2c[0], 0, scale2, -scale2*m2c[1],0,0,1])
    T2=T2.reshape([3,3])

    F=np.dot(T2.transpose(),np.dot(F,T1))

    #if np.abs(F[2,2]) >1.192092896e-07 :
    #    F=F/F[2,2]
    return F


def haveNearPoints(pts1,pts2,modelPoints):
    """
    分别检查view1,view2中是否有相互接近的点
    """
    for i in range(modelPoints - 1):
        for j in range(i):
            p1 = pts1[i];
            p2 = pts1[j]
            dist = np.linalg.norm(p1 - p2)
            if dist < 50:
                return True

    for i in range(modelPoints - 1):
        for j in range(i):
            p1 = pts2[i];
            p2 = pts2[j]
            dist = np.linalg.norm(p1 - p2)
            if dist < 50:
                return True
    return False


def haveCollinearPoints(pts1,pts2,modelPoints):
    FLT_EPSILON = 1.192092896e-07
    for i in range(modelPoints - 1):
        dx1 = pts1[i, 0] - pts1[modelPoints - 1, 0]  # 最后一个点做为固定点
        dy1 = pts1[i, 1] - pts1[modelPoints - 1, 1]
        for j in range(i):
            dx2 = pts1[j, 0] - pts1[modelPoints - 1, 0]
            dy2 = pts1[j, 1] - pts1[modelPoints - 1, 1]
            if (np.abs(dx2 * dy1 - dy2 * dx1)) <= FLT_EPSILON * (np.abs(dx1) + np.abs(dx2) + np.abs(dy1) + np.abs(dy2)):
                return True
    for i in range(modelPoints - 1):
        dx1 = pts2[i, 0] - pts2[modelPoints - 1, 0]  # 最后一个点做为固定点
        dy1 = pts2[i, 1] - pts2[modelPoints - 1, 1]
        for j in range(i):
            dx2 = pts2[j, 0] - pts2[modelPoints - 1, 0]
            dy2 = pts2[j, 1] - pts2[modelPoints - 1, 1]
            if (np.abs(dx2 * dy1 - dy2 * dx1)) <= FLT_EPSILON * (np.abs(dx1) + np.abs(dx2) + np.abs(dy1) + np.abs(dy2)):
                return True

    return False

def random_partition(n,N,pts1,pts2):
    """return n random rows of data (and also the other len(data)-n rows)"""
    """返回数据的随机n行的索引,以及剩余数据的索引"""
    while True:
        all_idxs = np.arange( N )
        np.random.shuffle(all_idxs)
        idxs = all_idxs[:n]
        #idxs2 = all_idxs[n:]
        ##check,两个点之间的距离不能太近
        if not haveNearPoints(pts1[idxs],pts2[idxs],n) and not haveCollinearPoints(pts1[idxs],pts2[idxs],n):
            break
    return idxs


class ransac_model_for_F(object):
    def __init__(self):
        return

    def fit(self,data1,data2,modelPoints):
        """从选择的7点对计算F矩阵 """
        #kp1=data[:,:2]
        #kp2=data[:,2:]
        Fs=[]
        if modelPoints ==7: #7点算法
            Fs=seven_point_algorithm(data1,data2)
        else:               #8点算法
            #modelPoints=8
            F=eight_point_algorithm(data1,data2)
            Fs.append(F)


        return Fs    #返加1到3个可行F矩阵

    def haveCollinearPoints(self,pts,modelPoints):
        FLT_EPSILON = 1.192092896e-07
        for i in range(modelPoints-1):
            dx1=pts[i,0]-pts[modelPoints-1,0]   #最后一个点做为固定点
            dy1=pts[i,1]-pts[modelPoints-1,1]
            for j in range(i):
                dx2=pts[j,0]-pts[modelPoints-1,0]
                dy2=pts[j,1]-pts[modelPoints-1,1]
                if(np.abs(dx2*dy1-dy2*dx1))<=FLT_EPSILON*(np.abs(dx1)+np.abs(dx2)+np.abs(dy1)+np.abs(dy2)):
                    return True
        return False

    def checkSamples(self,sd1,sd2,modelPoints):
        """对采样的匹配点对分别在两个图像上进行共线性检测"""
        if self.haveCollinearPoints(sd1,modelPoints) or self.haveCollinearPoints(sd2,modelPoints):
            return False
        elif self.haveNearPoints(sd1,modelPoints) or self.haveNearPoints(sd2,modelPoints):
            return False
        else: return True


def compute_distance(data1,data2,F):
    """计算点对列表中每个点对在模型下的误差
    data1,data2,Nx2特征匹配点
    F: 3x3基本矩阵
    """
    n=len(data1)
    dist=np.zeros((2,n))
    f=F.reshape(9)
    for i in range(n):
        x1=np.array([data1[i,0],data1[i,1],1])
        x2=np.array([data2[i,0],data2[i,1],1])
        #计算view1中特征点在view2中的极线
        line1=np.dot(F,x1)
        s=np.sqrt(line1[0]**2+line1[1]**2)
        line1=line1/s
        #计算其匹配点到极线的距离
        dist1=abs(np.dot(x2,line1))
        #计算view2中特征点在view1中的极线
        line2=np.dot(np.transpose(F),x2)
        s=np.sqrt(line2[0]**2+line2[1]**2)
        dist2=abs(np.dot(line2,x1))
        dist[0,i]=dist1
        dist[1,i]=dist2
    return dist

def get_correspondences(img1,img2):
    """
    返回值：    kp1--img1对应的匹配点坐标(x,y)tuple的list
                kp2--img2对应的匹配点坐标(x,y)tuple的list
    """
    ##step 1 detect keypoints & compute descriptor
    surf=cv2.xfeatures2d.SURF_create(4000)  #参数为Hessian阀值
    kp1, des1 = surf.detectAndCompute(img1,None)
    kp2, des2 = surf.detectAndCompute(img2,None)

    ##step 2 特征匹配
    FLANN_INDEX_KDTREE=1
    indexParams=dict(algorithm =FLANN_INDEX_KDTREE,trees=5)
    searchParams=dict(checks=50)
    flann=cv2.FlannBasedMatcher(indexParams,searchParams)
    #matches=flann.match(des1,des2)              #第一个作为quest desc，第二个作为train desc
    matches=flann.knnMatch(des1,des2,k=2)
    #matches=sorted(matches,key=lambda x:x.distance) #按距离排序

    goodMatches=[]; pts1=[]; pts2=[]
    for (m,n) in matches:
        if m.distance < 0.7*n.distance:
            goodMatches.append(m)
            pts2.append(kp2[m.trainIdx].pt)
            pts1.append(kp1[m.queryIdx].pt)
    """
    pts1=[]
    pts2=[]
    for x in matches:
        pts1.append(kp1[x.queryIdx].pt)
        pts2.append(kp2[x.trainIdx].pt)
    """
    """
    ######显示所有区配点
    img3=cv2.drawMatches(img1,kp1,img2,kp2,goodMatches,img2,flags=2)
    height,width=img3.shape[:2]
    img3=cv2.resize(img3,(np.int(width*0.6),np.int(height*0.6)))
    while(True):
        cv2.imshow('matches',img3)
        if cv2.waitKey() & 0xff==ord("q"):
            break
    cv2.destroyAllWindows()
    """

    #cv2.imwrite("e:\out.jpg",img3)
    return  np.array(pts1),np.array(pts2),np.array(goodMatches),kp1,kp2

def draw_epipolar_relation(img1,img2,F,e1,e2,p1,p2):
    """画出两幅图像中两个点的极线关系"""
    nWidth,nHeight=img1.shape[:2]
    ##img1中的p1点在img2中的极线及法方向
    cv2.circle(img2,tuple(e2[:2]),5,(0,255,0),thickness=2)
    line2=np.dot(F,p1)  #p1点在img2中对应的极线
    #显示这条极线
   # ep1=


    return


def computeF(img1,img2):
    """
    通过两幅图像计算F矩阵
    返回：F矩阵，极点e1,e2,有效匹配点对kp1,kp2
    """
    ##获得两幅图像的匹配
    pts1,pts2,goodMatches,kps1,kps2=get_correspondences(img1,img2)
    #pt1=np.int32(pt1)
    #pt2=np.int32(pt2)
    numOfDatas=len(pts1)

    ##调用RANSAC算法
    maxInliers=0
    bestResErr=np.inf
    bestInliers_idxs=[]
    e=0.66      #outlier probability
    p=0.99      # confidence
    thresh=1
    s=8         #s-points algorithm
    iter=0
    #N=np.log(1-p)/np.log(1-(1-e)**s)        #迭代次数
    N=1000
    while iter<N:
        iter=iter + 1
        fit_idxes =random_partition(s,numOfDatas,pts1,pts2)    #得到随机的s个点的索引
        F = eight_point_algorithm(pts1[fit_idxes], pts2[fit_idxes])     #运用8点算法
        ##计算几何误差
        dist=compute_distance(pts1,pts2,F)
        inlierCnt=np.sum(np.int8(dist[0] < thresh) * np.int8(dist[1] < thresh))
        inliers_idxs=np.where(np.int8(dist[0]<thresh)*np.int8(dist[1]<thresh))      #注意，多了一层tuple
        if inlierCnt>0:
            resErr=np.sum(dist[0,inliers_idxs]**2+dist[1,inliers_idxs]**2)/inlierCnt
        else:
            resErr=np.inf
        if inlierCnt>maxInliers or (inlierCnt==maxInliers and resErr<bestResErr):
            maxInliers=inlierCnt
            bestResErr=resErr
            bestF=F
            bestInliers_idxs=inliers_idxs[0]

            ##计算自适应迭代次数
            #e=1-inlierCnt/numOfDatas
            #if e>0:
            #    N=np.log(1-p)/np.log(1-(1-e)**s)
            #else:
            #    N=1

    if bestF is None:
        return None,None,None,None,None

    """
    #####显示符合的匹配点
    img3=cv2.drawMatches(img1,kps1,img2,kps2,goodMatches[bestInliers_idxs],img2,flags=2)
    height,width=img3.shape[:2]
    img3=cv2.resize(img3,(np.int(width*0.6),np.int(height*0.6)))
    while(True):
        cv2.imshow('matches',img3)
        if cv2.waitKey() & 0xff==ord("q"):
            break
    cv2.destroyAllWindows()
    print('*****F*****')
    print(np.linalg.det(bestF))
    print(bestF)
    """

    ##计算极点，e1,e2(e1为相机2中心C在img1中的坐标，e2为相机1中心C在img2中的坐标)
    Ft=bestF.transpose()
    U,s,V=np.linalg.svd(bestF)
    e1=V[-1,:]; e1=np.int32(e1/e1[2])
    U,s,V=np.linalg.svd(Ft)
    e2=V[-1,:]; e2=np.int32(e2/e2[2])
    print("*****polar point*****")
    print(e1)
    print(e2)
    return F,e1,e2,pts1[bestInliers_idxs],pts2[bestInliers_idxs]



"""test"""
img1 = cv2.imread('../misc_pic/apt1.jpg',0)
img2 = cv2.imread('../misc_pic/apt2.jpg',0)
#gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
#gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

N=100
e1a=np.zeros(N)
e1b=np.zeros(N)
e2a=np.zeros(N)
e2b=np.zeros(N)
for i in range(N):
    F,e1,e2,kp1,kp2=computeF(img1,img2)
    e1a[i]=e1[0]
    e1b[i]=e1[1]
    e2a[i]=e2[0]
    e2b[i]=e2[1]
    print(i)

plt.subplot(411)
plt.plot(e1a)
plt.subplot(412)
plt.plot(e1b)
plt.subplot(413)
plt.plot(e2a)
plt.subplot(414)
plt.plot(e2b)
plt.show()


"""
if F is not None:
    ##显示img1中点在img2中的极线,l=Fx
    width,height=img2.shape[:2]
    for kp in kp1[:10]:
        x=[kp[0],kp[1],1]
        line=np.dot(F,x)
        #normalize
        s=np.sqrt(line[0]**2+line[1]**2)
        line=line/s
        #print(line)
        #线上的点
        img2=mf.draw_line(img2,line,mf.rand_color())
    #plt.gray()
    plt.imshow(img2)  # 先显示img2
    plt.show()
    while(True):
        cv2.imshow('corners',img2)
        if cv2.waitKey() & 0xff==ord("q"):
            break
    cv2.destroyAllWindows()

"""
